#!/usr/bin/env python3

import sys
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QPushButton, QVBoxLayout, QHBoxLayout, QSpacerItem, QSizePolicy
import threading
import rclpy
import numpy as np
from sensor_msgs.msg import Image
from rclpy.node import Node
from std_msgs.msg import String
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import QTimer, Qt
from manipulation.msg import ManiProcessCmd, TargetsList
from cv_bridge import CvBridge
import cv2
import os
import time

os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH", None)

class ImageSubscriberNode(Node):
    def __init__(self):
        super().__init__('mani_gui_node')
        from cv_bridge import CvBridge
        self.bridge = CvBridge()
        self.cv_image = None
        self.targets_names = []
        self.targets_grasped = []  
        self.create_subscription(
            Image,
            '/detected_image',
            self.image_callback,
            10
        )
        self.create_subscription(
            TargetsList,
            '/detection_targets',
            self.targets_callback,
            10
        )
        self.mani_process_pub = self.create_publisher(
            ManiProcessCmd,
            '/mani_process2',
            10
        )

    def image_callback(self, msg):
        from cv_bridge import CvBridge
        import cv2
        try:
            self.cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as e:
            self.get_logger().error(f"cv_bridge exception: {e}")

    def targets_callback(self, msg):
        self.targets_names = list(msg.names)
        self.targets_grasped = list(msg.grasped)


class ManiWindow(QWidget):
    def __init__(self, ros_node: ImageSubscriberNode):
        super().__init__()
        from cv_bridge import CvBridge
        import cv2

        self.opentimes = 0

        current_time = time.localtime()
        date_str = time.strftime("%Y-%m-%d-%H-%M-%S", current_time)
        self.time = date_str

        self.ros_node = ros_node
        self.selected_index = 0  # 修改：默认选择第一个目标
        self.setWindowTitle("抓取操作界面")
        self.setGeometry(300, 300, 1500, 900)
        main_layout = QVBoxLayout(self)

        # 顶部布局，右上角显示"list"
        top_layout = QHBoxLayout()
        top_layout.addSpacerItem(QSpacerItem(120, 20, QSizePolicy.Expanding, QSizePolicy.Minimum))
        list_label = QLabel("list", self)
        list_label.setStyleSheet("font-size:18px;")
        top_layout.addWidget(list_label)       
        main_layout.addLayout(top_layout)

        # names显示区（紧跟list下方，右对齐）
        names_container = QWidget(self)
        names_hlayout = QHBoxLayout(names_container)
        names_hlayout.addSpacerItem(QSpacerItem(80, 20, QSizePolicy.Expanding, QSizePolicy.Minimum))
        self.names_vlayout = QVBoxLayout()
        self.names_labels = []
        names_hlayout.addLayout(self.names_vlayout)
        names_container.setLayout(names_hlayout)
        main_layout.addWidget(names_container)

        # 图像显示区（左侧）
        self.image_label = QLabel("等待图像...", self)
        self.image_label.setFixedSize(1200, 700)
        self.image_label.setAlignment(Qt.AlignLeft | Qt.AlignTop)
        main_layout.addWidget(self.image_label)
        main_layout.addStretch() 

        # 底部按钮布局，右下角，纵向排列
        bottom_row = QHBoxLayout()
        bottom_row.addStretch()
        btn_col = QVBoxLayout()
        self.prepare_btn = QPushButton("prepare", self)
        btn_col.addWidget(self.prepare_btn)
        self.start_btn = QPushButton("start", self)
        btn_col.addWidget(self.start_btn)
        self.place_btn = QPushButton("place origin", self)
        btn_col.addWidget(self.place_btn)
        self.reset_btn = QPushButton("reset", self)
        btn_col.addWidget(self.reset_btn)
        self.quit_btn = QPushButton("quit", self)
        btn_col.addWidget(self.quit_btn)
        btn_col.addStretch()
        bottom_row.addLayout(btn_col)
        main_layout.addLayout(bottom_row)

        self.setLayout(main_layout)

        self.prepare_btn.clicked.connect(self.on_prepare_clicked)
        self.start_btn.clicked.connect(self.on_start_clicked)
        self.place_btn.clicked.connect(self.on_place_clicked)
        self.reset_btn.clicked.connect(self.on_reset_clicked)
        self.quit_btn.clicked.connect(self.on_quit_clicked)

        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_image)
        self.timer.timeout.connect(self.update_names)
        self.timer.start(100)

        self.installEventFilter(self)

    def eventFilter(self, obj, event):
        from PyQt5.QtCore import QEvent
        if event.type() == QEvent.MouseButtonPress:
            pos = event.pos()
            in_names = False
            for label in self.names_labels:
                if label.geometry().contains(label.mapFromParent(pos)):
                    in_names = True
                    break
            if not in_names and self.selected_index is not None:
                self.selected_index = None
                self.update_names()
        return super().eventFilter(obj, event)
    
    def on_prepare_clicked(self):
        msg = ManiProcessCmd()
        msg.level1_name = 'prepare'
        self.ros_node.mani_process_pub.publish(msg)

    def on_start_clicked(self):
        import subprocess
        self.opentimes += 1
        if self.selected_index is not None:
            msg = ManiProcessCmd()
            msg.level1_name = 'start'
            msg.value = float(self.selected_index)
            self.ros_node.mani_process_pub.publish(msg)  # 抓取瓶子无需缺陷检测
            # cmd = (
            #     "source install/setup.sh && "
            #     f"ros2 run gui ros2_detection.py --ros-args -p date:='{self.time}' -p i:={self.opentimes}"
            # )
            # subprocess.Popen(["bash", "-c", cmd])

    def on_place_clicked(self):
        msg = ManiProcessCmd()
        msg.level1_name = 'place'
        msg.level2_name = 'origin'
        self.ros_node.mani_process_pub.publish(msg)

    def on_reset_clicked(self):
        msg = ManiProcessCmd()
        msg.level1_name = 'reset'
        self.ros_node.mani_process_pub.publish(msg)

    def on_quit_clicked(self):
        msg = ManiProcessCmd()
        msg.level1_name = 'quit'
        self.ros_node.mani_process_pub.publish(msg)
        self.close()
    
    def update_image(self):
        from cv_bridge import CvBridge
        import cv2
        if self.ros_node.cv_image is not None:
            rgb_image = cv2.cvtColor(self.ros_node.cv_image, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w
            qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
            self.image_label.setPixmap(QPixmap.fromImage(qt_image).scaled(
                self.image_label.width(), self.image_label.height(), Qt.KeepAspectRatio))
        else:
            self.image_label.setText("未收到图像")

    def update_names(self):
        # 清除旧的labels
        for label in self.names_labels:
            self.names_vlayout.removeWidget(label)
            label.deleteLater()
        self.names_labels.clear()
        
        # 添加新的labels
        names = self.ros_node.targets_names
        grasped = getattr(self.ros_node, "targets_grasped", [])
        
        # 修改：确保至少有一个目标时，selected_index为0
        if names and self.selected_index is None:
            self.selected_index = 0
            
        for i, name in enumerate(names):
            status = (
                "（已抓取）" if (i < len(grasped) and grasped[i] == True)
                else "（未抓取）"
            )
            label = QLabel(f"{name} {status}", self)
            label.setStyleSheet(
                "font-size:16px;" +
                ("background-color: green;" if self.selected_index == i else "")
            )
            label.setCursor(Qt.ArrowCursor)
            label.mousePressEvent = self.make_name_click_handler(i)
            self.names_vlayout.addWidget(label)
            self.names_labels.append(label)

    def make_name_click_handler(self, idx):
        def handler(event):
            self.selected_index = idx
            self.update_names()
        return handler

def ros_spin(node):
    rclpy.spin(node)

if __name__ == "__main__":
    rclpy.init()
    ros_node = ImageSubscriberNode()
    thread = threading.Thread(target=ros_spin, args=(ros_node,), daemon=True)
    thread.start()

    app = QApplication(sys.argv)
    window = ManiWindow(ros_node)
    window.show()
    sys.exit(app.exec_())
    ros_node.destroy_node()
    rclpy.shutdown()